'''
Example: 一只落在向日葵中间的蜜蜂

curl -d '{"chinese":"\u4e00\u53ea\u843d\u5728\u5411\u65e5\u8475\u4e2d\u95f4\u7684\u871c\u8702","seed":9}' http://127.0.0.1:8978

{"english": "A bee landed on the center of a sunflower.", "status": 0, "msg": "success"}
'''
import json
from http.server import BaseHTTPRequestHandler, HTTPServer
from ChatRWKV import ChatRWKV
import re
import os
import sys
import random
import numpy as np
import torch

FAKE_PHOTOS_RWKV_MODEL_PATH = os.environ.get('FAKE_PHOTOS_RWKV_MODEL_PATH')
if FAKE_PHOTOS_RWKV_MODEL_PATH is None:
    print('FAKE_PHOTOS_RWKV_MODEL_PATH is not set!', file=sys.stderr)
    exit(1)

FAKE_PHOTOS_RWKV_STRATEGY = os.environ.get('FAKE_PHOTOS_RWKV_STRATEGY') # cpu fp32, gpu fp16
if FAKE_PHOTOS_RWKV_STRATEGY is None:
    print('FAKE_PHOTOS_RWKV_STRATEGY is not set!', file=sys.stderr)
    exit(1)

RWKV = ChatRWKV(FAKE_PHOTOS_RWKV_MODEL_PATH, strategy=FAKE_PHOTOS_RWKV_STRATEGY)

def model(inp: str, rwkv: ChatRWKV) -> str:

    # 过滤所有换行符为空格
    prompt = re.sub(r'[\s|\n]+', ' ', inp)

    inputString = "Instruction: 按以下内容用英文描述画面\n\n"
    inputString = inputString + f"Input: {prompt}\n\n"
    inputString = inputString + "Response:"

    modelOutput = rwkv.run_on(
        inputString,
        temperature=1,
        top_p=0.3,
        top_k=15,
        max_length=800
    )
    rwkv.clear()
    return modelOutput

class MyHandler(BaseHTTPRequestHandler):

    def _handle_request(self, client_data: dict) -> str:
        if 'seed' not in client_data:
            seed = random.randint(0, 2**16 - 1)
        else:
            seed = client_data['seed']

        torch.manual_seed(seed)
        np.random.seed(seed)
        return model(client_data['chinese'], RWKV)

    def do_POST(self) -> None:
        # 获取POST推送过来的内容
        content_length = int(self.headers['Content-Length'])
        if content_length > 1024 * 1024:
            return

        post_data = self.rfile.read(content_length)
        try:
            client_data = json.loads(post_data.decode(encoding='utf8', errors='replace'))
            print(json.dumps(client_data, ensure_ascii=False))
        except json.decoder.JSONDecodeError as e:
            print(e)

        response = self._handle_request(client_data)

        self.send_response(200)
        self.send_header('Content-type', 'application/json')
        self.end_headers()
        self.wfile.write(bytes(json.dumps({'english':response,'status':0,'msg':'success'}), encoding='utf8'))

HTTPServer(('127.0.0.1', 8978), MyHandler).serve_forever()
